import torch
from torch import nn
from torch.nn import functional as F


class CenterLoss(nn.Module):
    def __init__(self, cls_num=2, feature_dim=(3, 224, 224)):
        super(CenterLoss, self).__init__()
        self.centers = nn.Parameter(torch.randn(cls_num, *feature_dim))

    def forward(self, features, labels, reduction='mean'):
        # 特征向量归一化
        _features = F.normalize(features)

        centers_batch = self.centers.index_select(dim=0, index=labels.long())
        # 根据论文《A Discriminative Feature Learning Approach for Deep Face Recognition》修改如下
        if reduction == 'sum':  # 返回loss的和
            return torch.sum(torch.pow(_features - centers_batch, 2)) / 2
        elif reduction == 'mean':  # 返回loss和的平均值，默认为mean方式
            return torch.sum(torch.pow(_features - centers_batch, 2)) / 2 / len(features)
        else:
            raise ValueError("ValueError: {0} is not a valid value for reduction".format(reduction))


class ArcFaceNet(nn.Module):
    def __init__(self, cls_num=10, feature_dim=2):
        super(ArcFaceNet, self).__init__()
        self.w = nn.Parameter(torch.randn(feature_dim, cls_num))

    def forward(self, features, m=1, s=10):
        # 特征与权重 归一化
        _features = F.normalize(features, dim=1)
        _w = F.normalize(self.w, dim=0)

        # 特征向量与参数向量的夹角theta，分子numerator，分母denominator
        theta = torch.acos(torch.matmul(_features, _w) / 10)  # /10防止下溢
        numerator = torch.exp(s * torch.cos(theta + m))
        denominator = torch.sum(torch.exp(s * torch.cos(theta)), dim=1, keepdim=True) - torch.exp(
            s * torch.cos(theta)) + numerator
        return torch.log(torch.div(numerator, denominator))
